// Tencent is pleased to support the open source community by making ncnn available.
//
// Copyright (C) 2025 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

#include <android/asset_manager_jni.h>
#include <android/native_window_jni.h>
#include <android/native_window.h>

#include <android/log.h>

#include <jni.h>

#include <string>
#include <vector>

#include <platform.h>
#include <benchmark.h>

#include "yolo11.h"


#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#if __ARM_NEON
#include <arm_neon.h>
#endif // __ARM_NEON

static int draw_unsupported(cv::Mat& rgb)
{
    const char text[] = "unsupported";

    int baseLine = 0;
    cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 1.0, 1, &baseLine);

    int y = (rgb.rows - label_size.height) / 2;
    int x = (rgb.cols - label_size.width) / 2;

    cv::rectangle(rgb, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)),
                    cv::Scalar(255, 255, 255), -1);

    cv::putText(rgb, text, cv::Point(x, y + label_size.height),
                cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 0, 0));

    return 0;
}

static int draw_fps(cv::Mat& rgb)
{
    // resolve moving average
    float avg_fps = 0.f;
    {
        static double t0 = 0.f;
        static float fps_history[10] = {0.f};

        double t1 = ncnn::get_current_time();
        if (t0 == 0.f)
        {
            t0 = t1;
            return 0;
        }

        float fps = 1000.f / (t1 - t0);
        t0 = t1;

        for (int i = 9; i >= 1; i--)
        {
            fps_history[i] = fps_history[i - 1];
        }
        fps_history[0] = fps;

        if (fps_history[9] == 0.f)
        {
            return 0;
        }

        for (int i = 0; i < 10; i++)
        {
            avg_fps += fps_history[i];
        }
        avg_fps /= 10.f;
    }

    char text[32];
    sprintf(text, "FPS=%.2f", avg_fps);

    int baseLine = 0;
    cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);

    int y = 0;
    int x = rgb.cols - label_size.width;

    cv::rectangle(rgb, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)),
                    cv::Scalar(255, 255, 255), -1);

    cv::putText(rgb, text, cv::Point(x, y + label_size.height),
                cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));

    return 0;
}

static YOLO11* g_yolo11 = 0;
static ncnn::Mutex lock;



extern "C" {

JNIEXPORT jint JNI_OnLoad(JavaVM* vm, void* reserved)
{
    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "JNI_OnLoad");



    ncnn::create_gpu_instance();

    return JNI_VERSION_1_4;
}

JNIEXPORT void JNI_OnUnload(JavaVM* vm, void* reserved)
{
    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "JNI_OnUnload");

    {
        ncnn::MutexLockGuard g(lock);

        delete g_yolo11;
        g_yolo11 = 0;
    }

    ncnn::destroy_gpu_instance();

}
//com/boby/homecamera/yolo/
// public native boolean loadModel(AssetManager mgr, int taskid, int modelid, int cpugpu);
JNIEXPORT jboolean JNICALL Java_com_boby_homecamera_yolo_YOLO11Ncnn_loadModel(JNIEnv* env, jobject thiz, jobject assetManager, jint taskid, jint modelid, jint cpugpu)
{
    if (taskid < 0 || taskid > 4 || modelid < 0 || modelid > 8 || cpugpu < 0 || cpugpu > 2)
    {
        return JNI_FALSE;
    }

    AAssetManager* mgr = AAssetManager_fromJava(env, assetManager);

    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "loadModel %p", mgr);

    const char* tasknames[5] =
    {
        "",
        "_seg",
        "_pose",
        "_cls",
        "_obb"
    };

    const char* modeltypes[9] =
    {
        "n",
        "s",
        "m",
        "n",
        "s",
        "m",
        "n",
        "s",
        "m"
    };

    std::string parampath = std::string("yolo11") + modeltypes[(int)modelid] + tasknames[(int)taskid] + ".ncnn.param";
    std::string modelpath = std::string("yolo11") + modeltypes[(int)modelid] + tasknames[(int)taskid] + ".ncnn.bin";
    bool use_gpu = (int)cpugpu == 1;
    bool use_turnip = (int)cpugpu == 2;

    // reload
    {
        ncnn::MutexLockGuard g(lock);

        {
            static int old_taskid = 0;
            static int old_modelid = 0;
            static int old_cpugpu = 0;
            if (taskid != old_taskid || (modelid % 3) != old_modelid || cpugpu != old_cpugpu)
            {
                // taskid or model or cpugpu changed
                delete g_yolo11;
                g_yolo11 = 0;
            }
            old_taskid = taskid;
            old_modelid = modelid % 3;
            old_cpugpu = cpugpu;

            ncnn::destroy_gpu_instance();

            if (use_turnip)
            {
                ncnn::create_gpu_instance("libvulkan_freedreno.so");
            }
            else if (use_gpu)
            {
                ncnn::create_gpu_instance();
            }

            if (!g_yolo11)
            {
                if (taskid == 0) g_yolo11 = new YOLO11_det;
                if (taskid == 1) g_yolo11 = new YOLO11_seg;
                if (taskid == 2) g_yolo11 = new YOLO11_pose;
                if (taskid == 3) g_yolo11 = new YOLO11_cls;
                if (taskid == 4) g_yolo11 = new YOLO11_obb;

                g_yolo11->load(mgr, parampath.c_str(), modelpath.c_str(), use_gpu || use_turnip);
            }
            int target_size = 320;
            if ((int)modelid >= 3)
                target_size = 480;
            if ((int)modelid >= 6)
                target_size = 640;
            g_yolo11->set_det_target_size(target_size);
        }
    }

    return JNI_TRUE;
}

// public native boolean openCamera(int facing);
JNIEXPORT jboolean JNICALL Java_com_boby_homecamera_yolo_YOLO11Ncnn_openCamera(JNIEnv* env, jobject thiz, jint facing)
{
    if (facing < 0 || facing > 1)
        return JNI_FALSE;

    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "openCamera %d", facing);



    return JNI_TRUE;
}

// public native boolean closeCamera();
JNIEXPORT jboolean JNICALL Java_com_boby_homecamera_yolo_YOLO11Ncnn_closeCamera(JNIEnv* env, jobject thiz)
{
    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "closeCamera");



    return JNI_TRUE;
}

// public native boolean setOutputWindow(Surface surface);
JNIEXPORT jboolean JNICALL Java_com_boby_homecamera_yolo_YOLO11Ncnn_setOutputWindow(JNIEnv* env, jobject thiz, jobject surface)
{
    ANativeWindow* win = ANativeWindow_fromSurface(env, surface);

    __android_log_print(ANDROID_LOG_DEBUG, "ncnn", "setOutputWindow %p", win);



    return JNI_TRUE;
}
JNIEXPORT jobjectArray JNICALL Java_com_boby_homecamera_yolo_YOLO11Ncnn_detectRGBA(
        JNIEnv* env, jobject thiz, jbyteArray rgba, jint width, jint height)
{
    if (!g_yolo11)
        return nullptr; // 没有模型加载

    // 获取 RGBA 数据
    jbyte* rgba_ptr = env->GetByteArrayElements(rgba, 0);

    // 转成 cv::Mat
    cv::Mat img(height, width, CV_8UC4, (unsigned char*)rgba_ptr);
    cv::Mat rgb;
    cv::cvtColor(img, rgb, cv::COLOR_RGBA2BGR);

    std::vector<Object> objects;
    {
        ncnn::MutexLockGuard g(lock);
        g_yolo11->detect(rgb, objects);
    }

    // 获取 Java 类和构造方法
    jclass objClass = env->FindClass("com/boby/homecamera/yolo/YOLO11Ncnn$Object");
    jmethodID ctor = env->GetMethodID(objClass, "<init>", "(IFFFFF)V");

    // 创建 Java 数组
    jobjectArray objArray = env->NewObjectArray(objects.size(), objClass, nullptr);

    for (size_t i = 0; i < objects.size(); i++)
    {
        const Object& o = objects[i];
        jobject jObj = env->NewObject(objClass, ctor, o.label, o.prob, o.rect.x, o.rect.y, o.rect.width, o.rect.height);
        env->SetObjectArrayElement(objArray, i, jObj);
    }

    env->ReleaseByteArrayElements(rgba, rgba_ptr, 0);

    return objArray;
}


JNIEXPORT jobjectArray JNICALL
Java_com_boby_homecamera_yolo_YOLO11Ncnn_detectRGBABuffer(JNIEnv* env, jobject thiz,
                                                  jobject rgbaBuffer, jint width, jint height, jint rotation)
{
    ncnn::MutexLockGuard g(lock);

    if (!g_yolo11)
    {
        return nullptr;
    }

    // 取 RGBA 指针
    unsigned char* rgba_ptr = (unsigned char*) env->GetDirectBufferAddress(rgbaBuffer);
    if (!rgba_ptr) return nullptr;

    // 构造 cv::Mat
    cv::Mat img(height, width, CV_8UC4, rgba_ptr);

    // 转成 RGB
    cv::Mat rgb;
    cv::cvtColor(img, rgb, cv::COLOR_RGBA2RGB);

    // 根据 rotation 调整方向
    if (rotation == 90)
        cv::rotate(rgb, rgb, cv::ROTATE_90_CLOCKWISE);
    else if (rotation == 180)
        cv::rotate(rgb, rgb, cv::ROTATE_180);
    else if (rotation == 270)
        cv::rotate(rgb, rgb, cv::ROTATE_90_COUNTERCLOCKWISE);

    // 调 YOLO 检测
    std::vector<Object> objects;
    g_yolo11->detect(rgb, objects);

    // 获取 Java 类和构造方法
    jclass objClass = env->FindClass("com/boby/homecamera/yolo/YOLO11Ncnn$Object");
    jmethodID ctor = env->GetMethodID(objClass, "<init>", "(IFFFFF)V");

    // 创建 Java 数组
    jobjectArray objArray = env->NewObjectArray(objects.size(), objClass, nullptr);

    for (size_t i = 0; i < objects.size(); i++)
    {
        const Object& o = objects[i];
        jobject jObj = env->NewObject(objClass, ctor, o.label, o.prob, o.rect.x, o.rect.y, o.rect.width, o.rect.height);
        env->SetObjectArrayElement(objArray, i, jObj);
    }


    return objArray;
}


}
